#!/usr/bin/env python
# -*- coding: utf-8 -*-
# @Time    : 2023/3/31 14:16
# @Author : 闫峰
# @File    : db_pool.py
import json
from datetime import datetime
import pymysql
from utils.readFileUtils.yamlControl import GetYamlData
from dbutils.pooled_db import PooledDB
from typing import Optional, Tuple, List, Dict, Union


class Database:
    def __init__(self, yaml_path, database_key, database):
        """
        yaml_path 数据库配置信息yaml 文件地址
        database_name  yaml文件中的数据库信息 key
        """
        self.mysql = GetYamlData(yaml_path).get_yaml_data()[database_key]
        self.pool = PooledDB(
            creator=lambda: pymysql.connect(
                host=self.mysql["host"],
                user=self.mysql["user"],
                password=self.mysql["password"],
                database=database,
                port=int(self.mysql["port"]),
                charset="utf8mb4",
                connect_timeout=9999,  # 设置超时时间
            ),
            maxconnections=100,
            ping=1,
        )
        self.conn = self.pool.connection()

    def execute(self, sql: str, value: Optional[Tuple] = None) -> int:
        if not isinstance(sql, str):
            raise ValueError("The sql parameter must be a string.")
        if value is not None and not isinstance(value, tuple):
            raise ValueError("The value parameter must be a tuple or None.")

        try:
            with self.conn.cursor() as cursor:
                if value is None:

                    rows = cursor.execute(sql)
                else:
                    rows = cursor.execute(sql, value)
                self.conn.commit()
                return rows
        except Exception as e:
            self.conn.rollback()
            print(f"{str(value)}更新失败")
            raise e
        finally:
            self.conn.close()

    def execute_many(self, sql: str, value: list) -> int:
        value = list(value)
        # 检查SQL语句和值的类型
        if not isinstance(sql, str):
            raise ValueError("The sql parameter must be a string.")
        if not isinstance(value, list):
            raise ValueError("The value parameter must be a list.")
        try:
            # 使用分批的方式执行插入，每批200条数据
            batch_size = 20
            num = 0
            for i in range(0, len(value), batch_size):
                batch_values = value[i:i + batch_size]
                with self.conn.cursor() as cursor:
                    rows = cursor.executemany(sql, batch_values)
                    num += rows
                self.conn.commit()
            return num  # 返回插入的总记录数
        except Exception as e:
            self.conn.rollback()
            print(sql)
            print(batch_values)
            print(e)
            raise e
        finally:
            self.conn.close()  # 确保连接最终被关闭

    def fetch_all(self, sql: str, value: Optional[Tuple] = None) -> List[Dict]:
        if not isinstance(sql, str):
            raise ValueError("The sql parameter must be a string.")
        if value is not None and not isinstance(value, tuple):
            raise ValueError("The value parameter must be a tuple or None.")
        result = []
        try:
            with self.conn.cursor(cursor=pymysql.cursors.DictCursor) as cursor:
                if not value:
                    cursor.execute(sql)
                else:
                    cursor.execute(sql, value)
                result = cursor.fetchall()
                # for res in result:
                #     self.verify(res)
                return result
        except Exception as e:
            print(sql)
            print(value)
            raise e
        finally:
            self.conn.close()

    def fetch_all_tuple(self, sql: str, value: Optional[Tuple] = None) -> List[Dict]:
        if not isinstance(sql, str):
            raise ValueError("The sql parameter must be a string.")
        if value is not None and not isinstance(value, tuple):
            raise ValueError("The value parameter must be a tuple or None.")
        try:
            with self.conn.cursor() as cursor:
                if not value:
                    cursor.execute(sql)
                else:
                    cursor.execute(sql, value)
                result = cursor.fetchall()
                # for res in result:
                #     self.verify(res)
                return result
        except Exception as e:
            print(sql)
            print(value)
            raise e
        finally:
            self.conn.close()

    def verify(self, result: dict) -> Union[dict, None]:
        """验证结果能否被json.dumps序列化"""
        # 尝试变成字符串，解决datetime 无法被json 序列化问题
        try:
            json.dumps(result)
        except TypeError:  # TypeError: Object of type datetime is not JSON serializable
            for k, v in result.items():
                if isinstance(v, datetime):
                    result[k] = str(v)[0:10]

    def fetch_max_id(self, sql: str) -> int:
        if not isinstance(sql, str):
            raise ValueError("The sql parameter must be a string.")
        result = []
        try:
            with self.conn.cursor(cursor=pymysql.cursors.DictCursor) as cursor:

                cursor.execute(sql)
                result = cursor.fetchall()

            return result
        except Exception as e:
            print(sql)
            raise e
